import bcrypt
from flask import jsonify
from database_connection import get_db_connection

def register_user(username, password):
    """
    注册用户
    :param username: 用户名
    :param password: 密码
    :return: JSON响应
    """
    cnx = get_db_connection()
    if cnx:
        try:
            cursor = cnx.cursor()
            # 检查用户名是否已经存在
            check_query = "SELECT * FROM users WHERE username = %s"
            data = (username,)
            cursor.execute(check_query, data)
            existing_user = cursor.fetchone()
            if existing_user:
                return jsonify({"error": "用户名已存在，请选择其他用户名"}), 400
            
            insert_query = "INSERT INTO users (username, password) VALUES (%s, %s)"
            hashed_password = bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt())
            data = (username, hashed_password.decode('utf-8'))
            cursor.execute(insert_query, data)
            cnx.commit()
            return jsonify({"message": "用户注册成功"})
        except Error as err:
            return jsonify({"error": f"注册用户时出错: {err}"}), 500
        finally:
            cursor.close()
            cnx.close()
    else:
        return jsonify({"error": "数据库连接失败"}), 500

def login_user(username, password):
    """
    登录用户
    :param username: 用户名
    :param password: 密码
    :return: JSON响应
    """
    cnx = get_db_connection()
    if cnx:
        try:
            cursor = cnx.cursor()
            select_query = "SELECT * FROM users WHERE username = %s"
            data = (username,)
            cursor.execute(select_query, data)
            user = cursor.fetchone()
            if user:
                stored_password = user[2].encode('utf-8')
                if bcrypt.checkpw(password.encode('utf-8'), stored_password):
                    return jsonify({"message": "登录成功"})
                else:
                    return jsonify({"error": "用户名或密码错误"}), 401
            else:
                return jsonify({"error": "用户名或密码错误"}), 401
        except Error as err:
            return jsonify({"error": f"登录时出错: {err}"}), 500
        finally:
            cursor.close()
            cnx.close()
    else:
        return jsonify({"error": "数据库连接失败"}), 500